% ***********************************************************************************************
% ***  m_counterfactualpath.m ******************************************************************* 
% ***********************************************************************************************
% compute different equilibrium when we have agglomeration in productivity and amenities 
% ***********************************************************************************************
clear all; 
close all;
format shortG;

% *************************************************************************
% ******************** LOAD FUNDAMENTAL PARAMETERS ************************
% *************************************************************************
load parameters.mat;                           
load calib_productivity_amenity.mat;    
load calib_setup.mat;                   
load commuting_cost.mat;  
load calib_auxiliary.mat;

load fundamentals_1950.mat;

load calib_qadjusted.mat; 
load calib_hscale.mat; 

load noagglomeration.mat; 

Rdata=Rdata_all; Ldata=Ldata_all;     
totalpop=sum(Rdata(:,2:end),1);

N=size(ID,1);
NT=6; 
% ************************************************************************
% ***************** EXOGENOUS FUNDAMENTALS (50-75) ***********************
% ************************************************************************
a_mat=[a50, aa];   a_mat=a_mat./repmat(geomean(a_mat),N,1);
b_mat=[b50, bb];   b_mat=b_mat./repmat(geomean(b_mat),N,1);
% ************************************************************************
% ******************** COMMUTING PATTERN 1945 ****************************
% ************************************************************************
Lcom=zeros(N,N,NT+1); 
Lcom(:,:,1)=Lcom45; 
R45=Rdata(:,1); L45=Ldata(:,1);

% *************************************************************************
% ************* INITIAL GUESS FOR COUNTERFACTUAL FROM 1950 ****************
% *************************************************************************
Q_i=ones(N,NT); 
L_i=(L45./sum(L45)).*totalpop; R_i=(R45./sum(R45)).*totalpop; 
A_i=a_mat.*(L_i./repmat(Area,1,NT)).^alphaest;  
B_i=b_mat.*(R_i./repmat(Area,1,NT)).^betaest;   

% *************************************************************************
% ***************** MAIN CALIBRATION STARTS *******************************
% *************************************************************************
Q_n=zeros(N,NT); L_n=zeros(N,NT); R_n=zeros(N,NT); 
H_n=zeros(N,NT+1); H_n(:,1)=hsupply45; H_demand=zeros(N,NT); 
maxit=1e+6; tol=1e-6; outupdate=0.25; update=0.05; it=0; 
while it<maxit 

% continuation value of floor spaces, productivity and amenities  
LQmat=zeros(N,NT);  Xmat=zeros(N,NT); Ymat=zeros(N,NT);   
LQmat(:,NT)=Q_i(:,NT); Xmat(:,NT)=B_i(:,NT); Ymat(:,NT)=A_i(:,NT);
for t=1:NT-1 
    tt=NT-t;
    theta=thetavec(tt); 
    LQmat(:,tt)=Q_i(:,tt).*(LQmat(:,tt+1).^(rho*(1-theta)));
    Xmat(:,tt)=B_i(:,tt).*(Xmat(:,tt+1).^(rho*(1-theta)));
    Ymat(:,tt)=A_i(:,tt).*(Ymat(:,tt+1).^(rho*(1-theta)));
end 

for t=1:NT
    if t==1; theta=theta50; Rpre_i=R45; Lpre_i=L45; 
    else     theta=thetavec(t-1); Rpre_i=R_n(:,t-1); Lpre_i=L_n(:,t-1); end 
    
    % conditional probabilities (lambda_L: residential given workplace; lambda_R: workplace given residenace) 
    lambda_L=Kmat(:,:,t).*((repmat(LQmat(:,t)',N,1).^(-mu)).*(repmat(Xmat(:,t)',N,1))).^(rho/sigma); 
    lambda_L=lambda_L./repmat(sum(lambda_L,2),1,N); 
    lambda_R=Kmat(:,:,t).*((repmat(LQmat(:,t),1,N).^(-gammaH/gammaL)).*(repmat(Ymat(:,t),1,N).^(1/gammaL))).^(rho/sigma);
    lambda_R=lambda_R./repmat(sum(lambda_R,1),N,1);
    
    % solve forward equations for population and employment 
    Rpost_i=R_i(:,t); Lpost_i=L_i(:,t);
    it1=0;
    while it1<maxit 
    Rpost_n=(1-theta).*Rpre_i+lambda_L'*(Lpost_i-(1-theta).*Lpre_i);
    Lpost_n=(1-theta).*Lpre_i+lambda_R*(Rpost_i-(1-theta).*Rpre_i);  
    if max(abs(Rpost_n-Rpost_i)) < tol && max(abs(Lpost_n-Lpost_i)) < tol; break;
    else
       Rpost_i=update.*Rpost_n+(1-update).*Rpost_i; 
       Lpost_i=update.*Lpost_n+(1-update).*Lpost_i;
       it1=it1+1;  
    end
    end
    R_n(:,t)=Rpost_n; L_n(:,t)=Lpost_n;

    % commuting patterns (column:=workplace; row:=residential place)
    Lcompost=(1-theta).*Lcom(:,:,t)+lambda_R.*repmat((Rpost_n-(1-theta).*Rpre_i)',N,1); 
    Lcom(:,:,t+1)=Lcompost;
    
    % wage (workplace) and average income (residential place) 
    w_n=(Q_i(:,t).^(-gammaH/gammaL)).*(A_i(:,t).^(1/gammaL));   
    v_n=(Lcompost./repmat(Rpost_n',N,1))'*w_n; 
    
    % aggregate demand and supply (stock) of floor space
    floor_demand=mu.*v_n.*Rpost_n + (gammaH/gammaL).*w_n.*Lpost_n;
    h_n=(1-psi).*H_n(:,t)+hbar(t)*(Q_i(:,t).^eta).*Area;
    H_n(:,t+1)=h_n; 
    H_demand(:,t)=floor_demand; 

    % floor space market clearing condition
    Q_n(:,t)=H_demand(:,t)./H_n(:,t+1) ;
end
 
A_n=a_mat.*(L_n./repmat(Area,1,NT)).^alphaest;  
B_n=b_mat.*(R_n./repmat(Area,1,NT)).^betaest;   
dd_Q=Q_n-Q_i; dd_A=A_n-A_i; dd_B=B_n-B_i; 

if max(max(abs(dd_A)))<tol && max(max(abs(dd_B)))<tol && max(max(abs(dd_Q)))<tol 
    disp(['converged population distribution']);
        break;
else
    Q_i=update.*Q_n+(1-update).*Q_i;
    B_i=update.*B_n+(1-update).*B_i; 
    A_i=update.*A_n+(1-update).*A_i; 
    it=it+1;  

   if rem(it,20)==1
        disp(['Largest Error is:', ' ',num2str(it), ' ',num2str(max(abs(dd_Q)))]);
   end
end

end

% check floor space per area, as well as the ratio of existing stock and
% total floor space
max(H_n./Area)

% check floor market clearing 
if max(abs(Q_n.*H_n(:,2:NT+1)-H_demand))<1e-5; disp(['PASS floor space market clearing conditions in 1950-75']);
else; disp(['WARNING']); end 

% *************************************************************************
% ****************************   SAVE RESULTS *****************************
% *************************************************************************
result=array2table([R_n,Rdata,L_n,Ldata, Q_n, Qadj, hstock, H_n, Lcom45(17,:)'./sum(Lcom45(17,:)), cbddist, Area, ID, lat, lon],'VariableNames',{'pop50_alt',...
    'pop55_alt','pop60_alt','pop65_alt','pop70_alt','pop75_alt',...
    'pop45','pop50','pop55','pop60','pop65','pop70','pop75',...
    'emp50_alt','emp55_alt','emp60_alt','emp65_alt','emp70_alt','emp75_alt',...
    'emp45','emp50','emp55','emp60','emp65','emp70','emp75',...
    'q50_alt','q55_alt','q60_alt','q65_alt','q70_alt','q75_alt',... 
    'q50','q55','q60','q65','q70','q75',... 
    'h50','h55','h60','h65','h70','h75',... 
    'h45','h50_alt','h55_alt','h60_alt','h65_alt','h70_alt','h75_alt',...
    'cbdcommute45','cbddist', 'area', 'geocode', 'lat', 'lon'}); 
writetable(result,'../../output/quant/output_counterfactualpath.csv');

Lcom_cf=Lcom; R_cf=R_n; L_cf=L_n; Q_cf=Q_n; A_cf=A_n; B_cf=B_n; 
save('counterfactualpath.mat','Lcom_cf','R_cf','L_cf','Q_cf','A_cf','B_cf');